clear
%constant
global gamma beta Bx Nx length dx dt eps
Bx = 5;
beta = 2;
gamma = 5/3;
length = 2;
time = 0.04;
eps = 1e-6;
% discretization
Nx = 261;
Nt = 400;
dx = length/(Nx-1);
dt = time/Nt;
u_array = zeros(7, Nx);
%initialization
for i = 1:1:Nx
    if i <= floor(Nx/2)
        u_array(1,i) = 3.108;   %rho
        u_array(2,i) = 1.4336;  %p
        u_array(3,i) = 0;       %vx
        u_array(4,i) = 0.2633;  %vy
        u_array(5,i) = 0.2633;  %vz
        u_array(6,i) = 0.1;     %By
        u_array(7,i) = 0.1;     %Bz
    else
        u_array(1,i) = 1.0;     %rho
        u_array(2,i) = 0.1;     %p
        u_array(3,i) = -0.9225; %vx
        u_array(4,i) = 0.0;     %vy
        u_array(5,i) = 0.0;     %vz
        u_array(6,i) = 1.0;     %By
        u_array(7,i) = 1.0;     %Bz   
    end
end

%time stepping
for i = 1:1:Nt
    u_array = u_array + time_step(u_array);
    i
end

rho_array = u_array(1,:);
p_array = u_array(2,:);
vx_array = u_array(3,:);
vy_array = u_array(4,:);
vz_array = u_array(5,:);
By_array = u_array(6,:);
Bz_array = u_array(7,:);

hold on
plot([-1:dx:1],rho_array);
plot([-1:dx:1],p_array);
plot([-1:dx:1],vx_array);
plot([-1:dx:1],vy_array);
plot([-1:dx:1],vz_array);
plot([-1:dx:1],By_array);
plot([-1:dx:1],Bz_array);

box on
set(gca,'XMinorTick','on')
set(gca,'YMinorTick','on')
xlabel('x')

function du_array = time_step(u_array)
    global Nx gamma eps length dx dt
    % calculate flux
    flux = Roe_flux(u_array);
    du_array = zeros(7,Nx);
    % refresh 
    for i = 2:1:Nx-1
        du_array(:,i) = -(dt/dx)*(flux(:,i) - flux(:,i-1));
    end
    % boundary points
    for i = 1:1:7
        du_array(i,1) = du_array(i,2);
        du_array(i,Nx) = du_array(i,Nx-1);
    end
end

% auxiliary parameters
function p = auxiliary(u)
% u: 1_rho 2_p 3_ux 4_uy 5_uz 6_By 7_Bz 
    global gamma Nx beta Bx
    rho = u(1);
    p = u(2);
    By = u(6);
    Bz = u(7);
    a = sqrt(beta*gamma*p/rho);
    bx = sqrt(Bx^2)/sqrt(rho);
    b_vert = sqrt(By^2+Bz^2)/sqrt(rho);
    b = sqrt(Bx^2+By^2+Bz^2)/sqrt(rho);
    cs = sqrt(0.5*(a^2+b^2-sqrt((a^2+b^2)^2-4*a^2*bx^2)));
    cf = sqrt(0.5*(a^2+b^2+sqrt((a^2+b^2)^2-4*a^2*bx^2)));
    eps_z = (Bz+eps)/sqrt(By^2+Bz^2+2*eps^2);
    eps_y = (By+eps)/sqrt(By^2+Bz^2+2*eps^2);
    alpha_s = sqrt((cf^2-a^2)/(cf^2-cs^2));
    alpha_f = sqrt((a^2-cs^2)/(cf^2-cs^2));

    p = [a;bx;b_vert;b;cs;cf;eps_y;eps_z;alpha_s;alpha_f];
end

function mat = Left_mat(u, para)
    global gamma beta Bx
    rho = u(1);
    vx = u(3);
    vy = u(4);
    vz = u(5);
    By = u(6);
    Bz = u(7);
    a = para(1);
    cs = para(5);
    cf = para(6);
    eps_y = para(7);
    eps_z = para(8);    
    alp_s = para(9);
    alp_f = para(10);
    s = sign(Bx);
    mat = [1, -beta/a^2, 0, 0, 0, 0, 0;
    0, 0, 0, eps_z/sqrt(2), -eps_y/sqrt(2), -eps_z*s/sqrt(2*rho), eps_y*s/sqrt(2*rho);
    0, 0, 0, eps_z/sqrt(2), -eps_y/sqrt(2), eps_z*s/sqrt(2*rho), -eps_y*s/sqrt(2*rho);
    0, beta*alp_s/a^2, rho*alp_s*cs/a^2, rho*alp_f*cf*eps_y*s/a^2, rho*alp_f*cf*eps_z*s/a^2, -sqrt(rho)*alp_f*eps_y/a, -sqrt(rho)*alp_f*eps_z/a;
    0, beta*alp_s/a^2, -rho*alp_s*cs/a^2, -rho*alp_f*cf*eps_y*s/a^2, -rho*alp_f*cf*eps_z*s/a^2, -sqrt(rho)*alp_f*eps_y/a, -sqrt(rho)*alp_f*eps_z/a;
    0, beta*alp_f/a^2, rho*alp_f*cf/a^2, -rho*alp_s*cs*eps_y*s/a^2, -rho*alp_s*cs*eps_z*s/a^2, sqrt(rho)*alp_s*eps_y/a, sqrt(rho)*alp_s*eps_z/a;
    0, beta*alp_f/a^2, -rho*alp_f*cf/a^2, rho*alp_s*cs*eps_y*s/a^2, rho*alp_s*cs*eps_z*s/a^2, sqrt(rho)*alp_s*eps_y/a, sqrt(rho)*alp_s*eps_z/a
    ]; 

    coef2 = (gamma-1)/beta;
    T = [1, 0, 0, 0, 0, 0, 0;
    coef2*0.5*(vx^2+vy^2+vz^2), coef2, -coef2*vx, -coef2*vy, -coef2*vz, -coef2*By, -coef2*Bz;
    -vx/rho, 0, 1/rho, 0, 0 ,0 ,0;
    -vy/rho, 0, 0, 1, 0, 0, 0;
    -vz/rho, 0, 0, 0, 1, 0, 0;
    0, 0, 0, 0, 0, 1, 0;
    0, 0, 0, 0, 0, 0, 1];
    
    mat = mat;%*T;
end

function mat = Right_mat(u, para)
    global gamma beta Bx
    rho = u(1);
    vx = u(3);
    vy = u(4);
    vz = u(5);
    By = u(6);
    Bz = u(7);
    a = para(1);
    cs = para(5);
    cf = para(6);
    eps_y = para(7);
    eps_z = para(8);    
    alp_s = para(9);
    alp_f = para(10);
    s = sign(Bx);
    mat = [1, 0, 0, 0, 0, 0, 0;
    0, 0, 0, eps_z/sqrt(2), -eps_y/sqrt(2), -eps_z*s*sqrt(rho/2), eps_y*s*sqrt(rho/2);
    0, 0, 0, eps_z/sqrt(2), -eps_y/sqrt(2), eps_z*s*sqrt(rho/2), -eps_y*s*sqrt(rho/2);
    alp_s/2, a^2*alp_s/(2*beta), alp_s*cs/(2*rho), alp_f*cf*eps_y*s/(2*rho), alp_f*cf*eps_z*s/(2*rho), -a*alp_f*eps_y/(2*sqrt(rho)), -a*alp_f*eps_z/(2*sqrt(rho));
    alp_s/2, a^2*alp_s/(2*beta), -alp_s*cs/(2*rho), -alp_f*cf*eps_y*s/(2*rho), -alp_f*cf*eps_z*s/(2*rho), -a*alp_f*eps_y/(2*sqrt(rho)), -a*alp_f*eps_z/(2*sqrt(rho));
    alp_f/2, a^2*alp_f/(2*beta), alp_f*cf/(2*rho), -alp_s*cs*eps_y*s/(2*rho), -alp_s*cs*eps_z*s/(2*rho), a*alp_s*eps_y/(2*sqrt(rho)), a*alp_s*eps_z/(2*sqrt(rho));
    alp_f/2, a^2*alp_f/(2*beta), -alp_f*cf/(2*rho), alp_s*cs*eps_y*s/(2*rho), alp_s*cs*eps_z*s/(2*rho), a*alp_s*eps_y/(2*sqrt(rho)), a*alp_s*eps_z/(2*sqrt(rho))
    ]; 

    T = [1, 0, 0, 0, 0, 0, 0;
    0.5*(vx^2+vy^2+vz^2), beta/(gamma-1), rho*vx, rho*vy, rho*vz, By, Bz;
    vx, 0, rho, 0, 0, 0, 0;
    vy, 0, 0, rho, 0, 0, 0;
    vz, 0, 0, 0, rho, 0, 0;
    0, 0, 0, 0, 0, 1, 0;
    0, 0, 0, 0, 0, 0, 1];

    mat = transpose(mat);
%    mat = T*transpose(mat);
end

function flux = Roe_flux(u_array)
    global gamma beta Bx Nx dt dx
    flux = zeros(7, Nx-1);

    fl = zeros(7,Nx-1);
    alpha_array = zeros(7,Nx-1);
    lambda_array = zeros(7,Nx-1);
    Right_array = zeros(7,7,Nx-1);
    gt_array = zeros(7,Nx-1);
    g_array = zeros(7, Nx-1);
    nu = zeros(7,Nx-1);
    for i = 2:1:Nx-1
        uL = u_array(:,i);
        uR = u_array(:,i+1);
        rho_L = sqrt(uL(1));
        rho_R = sqrt(uR(1));
        rho_avg = rho_L*rho_R;
        ux_avg = (rho_L*uL(3)+rho_R*uR(3))/(rho_R+rho_L);
        uy_avg = (rho_L*uL(4)+rho_R*uR(4))/(rho_R+rho_L);
        uz_avg = (rho_L*uL(5)+rho_R*uR(5))/(rho_R+rho_L);
        By_avg = (rho_R*uL(6)+rho_L*uR(6))/(rho_R+rho_L);
        Bz_avg = (rho_R*uL(7)+rho_L*uR(7))/(rho_R+rho_L);
        HL = 0.5*(uL(3)^2+uL(4)^2+uL(5)^2) + (uL(6)^2+uL(7)^2)/uL(1) + gamma*uL(2)/(gamma-1)/uL(1);
        HR = 0.5*(uR(3)^2+uR(4)^2+uR(5)^2) + (uR(6)^2+uR(7)^2)/uR(1) + gamma*uR(2)/(gamma-1)/uR(1);
        H_avg = (rho_L*HL + rho_R*HR)/(rho_R+rho_L);
        p_avg = (H_avg-0.5*(ux_avg^2+uy_avg^2+uz_avg^2) - (By_avg^2+Bz_avg^2)/rho_avg)*(gamma-1)*rho_avg/gamma;
        u_avg = [rho_avg, p_avg, ux_avg, uy_avg, uz_avg, By_avg, Bz_avg];
        para_avg = auxiliary(u_avg);
        vx = u_avg(3);
        bx = para_avg(2);
        cs = para_avg(5);
        cf = para_avg(6);
        lambda = [vx; vx+bx; vx-bx; vx+cs; vx-cs; vx+cf; vx-cf];
        alpha = Left_mat(u_avg,para_avg)*(uR-uL);
        R = Right_mat(u_avg,para_avg);
        alpha_array(:,i) = alpha;
        lambda_array(:,i) = lambda;  
        Right_array(:,:,i) = R;

        for k = 1:1:7
            nu(k,i) = dt/dx * lambda(k);
            g_tlide = 0.5*(Q(nu(k,i)) - nu(k,i)^2)*alpha(k);
            s = sign(g_tlide);      
            g = s*max(0, min(abs(g_tlide),gt_array(k,i-1)*s));
            gt_array(k,i) = g_tlide;        
            g_array(k,i) = g;
        end                         
    end

    for i = 2:1:Nx-2
        for k = 1:1:7
            if(alpha_array(k,i) == 0)
                gam(k,i) = 0;
            else
                gam(k,i) = (g_array(k,i+1)-g_array(k,i))/alpha_array(k,i);
            end
            coef = (g_array(k,i)+g_array(k,i+1)-Q(nu(k,i)+gam(k,i))*alpha_array(k,i));
            fl(:,i) = fl(:,i) + Right_array(:,k,i)*coef;
        end
    end

    for i = 1:1:Nx-1
        uL = u_array(:,i);
        par_L = auxiliary(uL);
        vx = uL(3);
        bx = par_L(2);
        cs = par_L(5);
        cf = par_L(6);
        flux_L = Right_mat(uL,par_L)*diag([vx; vx+bx; vx-bx; vx+cs; vx-cs; vx+cf; vx-cf])*Left_mat(uL,par_L)*uL;

        uR = u_array(:,i+1);
        par_R = auxiliary(uR);
        vx = uR(3);
        bx = par_R(2);
        cs = par_R(5);
        cf = par_R(6);
        flux_R = Right_mat(uR,par_R)*diag([vx; vx+bx; vx-bx; vx+cs; vx-cs; vx+cf; vx-cf])*Left_mat(uR,par_R)*uR;

        flux(:,i) = (flux_L + flux_R)/2;
    end
    flux(:,i) = flux(:,i) + fl(:,i);
end

function y = Q(x)
    if abs(x) > 0.2
        y = x^2/0.4 + 0.1;
    else
        y = abs(x);
    end
end

